import torch
from sklearn.metrics import confusion_matrix
from torch import nn
from torch.utils.data import DataLoader

from MyDataset import WindIcingDatasetV1, MyDataloader, WindIcingDatasetV2, SelectedFeatures, WindIcingDatasetV3
from MyTrainer import WindIcingTrainerV2, WindIcingTrainerV3
from MyUtil import try_gpu, get_logger, metrics_function_v1, get_best_file
from MyModel import WindIcingModelV1, WindIcingModelV2, WindIcingModelV3
from MyTester import WindIcingTesterV1
from my_train import no, get_model, get_dataset, get_trainer

if __name__ == '__main__':
    no = no
    logger = get_logger(no)

    
    checkpoint_path = get_best_file(no)
    logger.info(f'load checkpoint from {checkpoint_path}')
    checkpoint = torch.load(checkpoint_path, map_location=torch.device('cpu'))
    
    config = checkpoint['config']

    device = try_gpu(config["gpu_no"])
    logger.info(f'testing on device: {device}')
    
    
    get_dataset(config)

    model = get_model(config)

    model.load_state_dict(checkpoint['state_dict'])

    trainer = get_trainer(model, config)
    

    trainer.test(test_loader=config["data_dict"]["test_loader"])
    
    
